package com.ruoyi.ade.llm.agent;


import com.alibaba.fastjson.JSONObject;
import com.google.common.base.Strings;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.ruoyi.ade.llm.model.AdeMilvusRequest;
import com.ruoyi.ade.llm.model.AdeMilvusResponse;
import com.ruoyi.common.utils.StringUtils;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.vector.request.GetReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.GetResp;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;


/**
 * milvus客户端
 */
public class MilvusAgent {

    //http://192.168.254.201:9091/webui/
    private final String uri = "http://192.168.254.201:19530";
    private final int dimension = 768;
    private MilvusClientV2 client;

    //MIlvus中集合的schema列
    private final String C_ID = "id";
    private final String C_VECTOR = "vector";
    private final String C_TEXT = "text";
    private final String C_ADE = "ade";
    private final String C_EXPLAIN = "explain";


    private MilvusAgent() {
        ConnectConfig config = ConnectConfig.builder().uri(uri).build();
        client = new MilvusClientV2(config);
    }

    // 静态内部类负责实例化，类加载时是延迟的
    private static class SingletonHelper {
        private static final MilvusAgent INSTANCE = new MilvusAgent();
    }

    // 提供全局访问点
    public static MilvusAgent getInstance() {
        return SingletonHelper.INSTANCE;
    }

    /**
     * 创建集合
     *
     * @param collection
     */
    public void createCollection(String collection) {

        //检测集合是否存在
        HasCollectionReq hasCollectionReq = HasCollectionReq.builder()
                .collectionName(collection)
                .build();
        boolean has = client.hasCollection(hasCollectionReq);
        System.out.println("当前collection存在检测:" + has);
        if (has) return;

        CreateCollectionReq.CollectionSchema schema = client.createSchema();
        //创建字段
        schema.addField(AddFieldReq.builder().fieldName(C_ID).dataType(DataType.VarChar).isPrimaryKey(true).build());
        schema.addField(AddFieldReq.builder().fieldName(C_VECTOR).dataType(DataType.FloatVector).dimension(dimension).build());
        schema.addField(AddFieldReq.builder().fieldName(C_TEXT).dataType(DataType.VarChar).build());
        schema.addField(AddFieldReq.builder().fieldName(C_ADE).dataType(DataType.VarChar).build());
        schema.addField(AddFieldReq.builder().fieldName(C_EXPLAIN).dataType(DataType.VarChar).build());
        //创建索引
        IndexParam indexParam = IndexParam.builder()
                .fieldName("vector")
                .metricType(IndexParam.MetricType.COSINE)  //余弦相似度
                .build();
        //创建索引
        CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
                .collectionName(collection)
                .collectionSchema(schema)
                .indexParams(Collections.singletonList(indexParam))
                .build();
        client.createCollection(createCollectionReq);
        System.out.println("collection:" + collection + "创建成功了!");
    }


    /**
     * 索引数据
     */
    public void index(String collection, AdeMilvusRequest request) {
        //生成uuid
        if (Strings.isNullOrEmpty(request.getId()))
            throw new RuntimeException("AdeMilvusRequest: id can not be empty!");
        if (request.getText() == null)
            throw new RuntimeException("AdeMilvusRequest: text can not be empty!");
        if (request.getAde() == null)
            throw new RuntimeException("AdeMilvusRequest: ade can not be empty!");
        if (request.getExplain() == null)
            throw new RuntimeException("AdeMilvusRequest: explain can not be empty!");
        //创建插入对象
        JsonObject obj = new JsonObject();
        Gson gson = new Gson();
        List<Float> vectorList = Arrays.asList(request.getVector());
        obj.add(C_VECTOR, gson.toJsonTree(vectorList));
        obj.addProperty(C_ID, request.getId());
        obj.addProperty(C_TEXT, request.getText());
        obj.addProperty(C_ADE, request.getAde());
        obj.addProperty(C_EXPLAIN, request.getExplain());
        //插入对象
        InsertReq insertReq = InsertReq.builder()
                .collectionName(collection)
                .data(Collections.singletonList(obj))
                .build();
        client.insert(insertReq);
//        System.out.println("插入成功:" + request.getText());
    }

    /**
     * 单挑记录查询
     *
     * @param id
     * @param collection
     * @return
     */
    public GetResp get(String id, String collection) {
        GetReq req = GetReq.builder()
                .collectionName(collection)
                .ids(Collections.singletonList(id))
                .build();
        GetResp resp = client.get(req);
        List<QueryResp.QueryResult> list = resp.getResults;
        list.forEach(e -> {
            System.out.println(JSONObject.toJSONString(e));
        });
        return resp;
    }


    /**
     * @param collection 需要查询的集合
     * @param vec        需要查询的多个向量
     * @param score      相关性分数阈值
     * @param topN       top-n    每个检索句子都返回10条，如果10个检索句子 可能返回100条
     * @return 返回的结果已经去重了
     */
    public List<AdeMilvusResponse> search(String collection, Float[][] vec, float score, int topN) {

        List<FloatVec> list = new ArrayList<>();
        for (int i = 0; i < vec.length; i++) {
            List<Float> vectorList = Arrays.asList(vec[i]);
            list.add(new FloatVec(vectorList));
        }
        //构建检索请求
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collection)
                .data(new ArrayList<>(list))
                .outputFields(Collections.singletonList("*"))
                .topK(topN)
                .build();
        SearchResp resp = client.search(searchReq);
        List<List<SearchResp.SearchResult>> searchResults = resp.getSearchResults();
        List<AdeMilvusResponse> responses = new ArrayList<>();
        String id, text, ade, explain;
        for (List<SearchResp.SearchResult> results : searchResults) {
            for (SearchResp.SearchResult r : results) {
                if (r.getScore() < score) continue;
                id = r.getId().toString();
                text = r.getEntity().get(C_TEXT).toString();
                ade = r.getEntity().get(C_ADE).toString();
                explain = r.getEntity().get(C_EXPLAIN).toString();
                AdeMilvusResponse obj = new AdeMilvusResponse(id, text, ade, explain, r.getScore());
                responses.add(obj);
            }
        }
        return responses;
    }

    /**
     * @param collection 需要检索的集合
     * @param vector     需要查询的向量
     * @param score      相似性分数阈值
     * @param topN       top-n
     * @return
     */
    public List<AdeMilvusResponse> search(String collection, Float[] vector, float score, int topN) {

        List<Float> vectorList = Arrays.asList(vector);
        //构建检索请求
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collection)
                .data(Collections.singletonList(new FloatVec(vectorList)))
                .outputFields(Collections.singletonList("*"))
                .topK(topN)
                .build();
        SearchResp resp = client.search(searchReq);
        List<List<SearchResp.SearchResult>> searchResults = resp.getSearchResults();
        List<SearchResp.SearchResult> results = searchResults.get(0);
        List<AdeMilvusResponse> list = new ArrayList<>();
        String id, text, ade, explain;
        for (SearchResp.SearchResult r : results) {
            if (r.getScore() < score) continue;
            id = r.getId().toString();
            text = r.getEntity().get(C_TEXT).toString();
            ade = r.getEntity().get(C_ADE).toString();
            explain = r.getEntity().get(C_EXPLAIN).toString();
            AdeMilvusResponse obj = new AdeMilvusResponse(id, text, ade, explain, r.getScore());
            list.add(obj);
        }
        return list;
    }

    /**
     * 删除collection
     *
     * @param collection
     */
    public boolean dropCollection(String collection) {

        if(collection == null)  return true ;
            DropCollectionReq dropCollectionReq = DropCollectionReq.builder()
                    .collectionName(collection)
                    .build();
            client.dropCollection(dropCollectionReq);
        // check if dropped
        return !client.hasCollection(HasCollectionReq.builder()
                .collectionName(collection)
                .build());
    }

}
